from lib.base import BaseMetric
import torch
from torch.nn import functional as F


class MamlAcc(BaseMetric):
    def __init__(self, cfg):
        super().__init__(cfg)
        self.cfg = cfg

    def __call__(self, logits_q, labels, type):
        # logits_q: 3000*dim

        batch_size = self.batch_size[type]
        n_way = self.n_way[type]
        n_query = self.n_query[type]

        labels = torch.arange(n_way).to(labels.device)
        # labels: 60 -> 1*60*1 -> 10*60*5
        labels = labels.view(1, n_way, 1).expand(batch_size, n_way, n_query)
        # labels: 10*60*5 -> 3000
        labels = labels.flatten().to(torch.int64)
        # pred_q: 3000*1
        pred_q = F.softmax(logits_q, dim=1).argmax(dim=1)
        correct = torch.eq(pred_q, labels).sum().item()

        return correct / len(pred_q)
